import os
import sys
import glob
import time
import json
from sentence_transformers import SentenceTransformer
from elasticsearch import Elasticsearch
from elasticsearch.helpers import parallel_bulk
from PIL import Image
from tqdm import tqdm
from datetime import datetime
from exif import Image as exifImage


"""
glob：内置模块，该模块主要是用来查找文件与目录的
PIL：图像处理
tqdm：进度条
exif：图片元数据
SentenceTransformers：用于句子、文本和图像嵌入的Python库。为 100 多种语言计算文本的嵌入并且可以轻松地将它们用于语义文本相似性、语义搜索和同义词挖掘等常见任务。 该框架基于 PyTorch 和 Transformers，并提供了大量针对各种任务的预训练模型。
"""


def main():
    print(f'\n第一步：开始加载模型')
    start_time = time.perf_counter()
    # img_model = SentenceTransformer('clip-ViT-B-32')
    img_model = SentenceTransformer('/Users/lrk/.cache/huggingface/hub/clip-ViT-B-32')  # 加载本地库
    duration = time.perf_counter() - start_time
    print(f'加载模型耗时 = {duration}')

    print("\n第二步：创建图像嵌入")
    doc_list = []
    chunk_size = 100
    start_time = time.perf_counter()

    for filename in tqdm(file_paths, desc='正在处理文件', total=len(file_paths)):
        img = Image.open(filename)
        embedding = img_model.encode(img)
        name = os.path.basename(filename)
        doc = {
            'image_id': name.split(".")[0],    # 文件ID
            'image_name': name,                # 文件名
            'image_embedding': embedding.tolist(),       # 图片向量
            'relative_path': os.path.relpath(filename),  # 相对路径
            'exif': {}                                   # EXIF 信息
        }
        exif_data = get_exif_date(filename) # 日期信息
        if exif_data:
            doc['exif']['date'] = exif_data
        exif_location = get_exif_location(filename) # 照片GPS位置
        if exif_location:
            doc['exif']['location'] = exif_location
        doc_list.append(doc)

    if not doc_list:
        raise Exception("没有找到任何图像文件.")

    duration = time.perf_counter() - start_time
    print(f'创建图像嵌入耗时 = {duration}')

    print("\n第三步：创建 ES 索引写入向量")
    with open("image-embeddings-mappings.json", "r") as config_file:
        config = json.loads(config_file.read())
        if es.indices.exists(index=index_name):
            print("删除现有的索引 %s" % index_name)
            es.indices.delete(index=index_name)

        print("创建索引 %s" % index_name)
        es.indices.create(index=index_name, mappings=config["mappings"], settings=config["settings"])

    count = 0
    for success, info in parallel_bulk(index=index_name, client=es, actions=doc_list, thread_count=4, chunk_size=chunk_size, timeout='%ss' % 120):
        if success:
            count += 1
            if count % chunk_size == 0:
                print('Indexed %s documents' % str(count), flush=True)
                sys.stdout.flush()
        else:
            print('Doc 失败', info)

    print('写入文档数 %s ' % str(count), flush=True)
    duration = time.perf_counter() - start_time
    print(f'总耗时 = {duration}')
    print("完成!\n")


def get_exif_date(filename):
    try:
        with open(filename, 'rb') as f:
            image = exifImage(f)
            taken = f"{image.datetime_original}"
            date_object = datetime.strptime(taken, "%Y:%m:%d %H:%M:%S")
            return date_object.isoformat()
    except Exception as e:
        pass


def get_exif_location(filename):
    try:
        with open(filename, 'rb') as f:
            img = exifImage(f)
            if img.gps_latitude and img.gps_latitude_ref:
                lat = dms_coordinates_to_dd_coordinates(img.gps_latitude, img.gps_latitude_ref)
                lon = dms_coordinates_to_dd_coordinates(img.gps_longitude, img.gps_longitude_ref)
                return [lon, lat]
    except Exception as e:
        pass


def dms_coordinates_to_dd_coordinates(coordinates, coordinates_ref):
    decimal_degrees = coordinates[0] + coordinates[1] / 60 + coordinates[2] / 3600
    if coordinates_ref == "S" or coordinates_ref == "W":
        decimal_degrees = -decimal_degrees
    return decimal_degrees


if __name__ == '__main__':
    # 图片目录地址
    # PATH_TO_IMAGES = './images/**/*.jp*g'
    file_paths = glob.glob('../static/images/animal/**/*.jp*g', recursive=True)
    file_paths += glob.glob('../static/images/animal/**/*.png', recursive=True)

    # 创建 es 连接
    es = Elasticsearch('https://elastic:123456@localhost:9200', ca_certs='../certs/ca.crt', verify_certs=True, request_timeout=3600)
    index_name = 'my-image-embeddings'

    main()
